"""
from multiview_calib.calibpkl_predict import CalibPredict
calibPredict = CalibPredict(calibfile)
p2d = calibPredict.p3d_to_p2d(p3d)
p3d = calibPredict.p2d_to_p3d(p2d)
"""
# %% imports
from builtins import NotImplementedError, isinstance
import numpy as np
import pickle
import os
import itertools
import cv2


class CalibPredict:
    def __init__(self, calibfile):
        """
        init by the 'calibfile' as pkl file, or 'calib' as dict.
        """
        if isinstance(calibfile, str):
            assert os.path.isfile(calibfile)
            with open(calibfile, 'rb') as f:
                calib = pickle.load(f)
        elif isinstance(calibfile, dict):
            calib = calibfile
        else:
            raise NotImplementedError
        assert {'ba_poses'} <= set(calib.keys()), 'calib file is not correct'
        self.poses = calib['ba_poses']
        self.image_shape = list(calib['intrinsics'].values())[0]['image_shape'] if 'intrinsics' in calib else None #HxW
        self.views = sorted(list(calib['ba_poses'].keys()))
        for view in self.views:
            for item in ['K', 'R', 't', 'dist']:
                self.poses[view][item] = np.array(self.poses[view][item])


    def p3d_to_p2d(self, p3d:np.ndarray, image_shape=None) -> np.ndarray:
        """
        Convert 3D points to 2D points.
        ---------- inputs ----------
        p3d: nsample_**_xyz
        image_shape: (H, W)

        ---------- outputs ----------
        p2d: nview_nsample_**_xy
        """
        assert p3d.shape[-1] == 3
        p3d_flatten = p3d.astype(float).reshape(-1, 3)
        if image_shape is None and self.image_shape is not None:
            image_shape = self.image_shape
        nviews = len(self.poses)
        p2d = np.zeros((nviews, p3d_flatten.shape[0], 2), dtype=np.float64) + np.nan   # nviews_nsample_2
        views = self.views
        if image_shape is not None:
            mask_inside = lambda proj: np.logical_and.reduce([proj[:,0]>0, proj[:,0]<image_shape[1],
                                                proj[:,1]>0, proj[:,1]<image_shape[0]])
                    
        for view in views:
            param = self.poses[view]
            K,R,t,dist = param['K'], param['R'], param['t'], param['dist']
            p2d_tmp, _ = project_points(p3d_flatten, K, R, t, dist, image_shape)
            if image_shape is not None:
                p2d_tmp[~mask_inside(p2d_tmp)] = np.nan
            p2d[view] = p2d_tmp
        p2d = p2d.reshape((nviews, *p3d.shape[:-1], 2))
        return p2d

    def p2d_to_p3d(self, p2d:np.ndarray) -> np.ndarray:
        """
        Convert 2D points to 3D points.
        ---------- inputs ----------
        p2d: nview_nsample_**_xy
        image_shape: (H, W)

        ---------- outputs ----------
        p3d: nsample_**_xyz
        """
        original_shape = p2d.shape
        views = self.views
        nviews = len(views)
        assert original_shape[-1]==2
        assert nviews == original_shape[0]
        landmarks = p2d.reshape(original_shape[0], -1, original_shape[-1])
        p3d = build_input_np(views, self.poses, landmarks)
        p3d = p3d.reshape(*original_shape[1:-1], 3)
        return p3d

    def p3d_alignby_cam(self, icam:int, p3d:np.ndarray) -> np.ndarray:
        """
        Convert 3D points to 3D points.
        ---------- inputs ----------
        p3d: nsample_**_xy
        image_shape: (H, W)

        ---------- outputs ----------
        p3d_align: nsample_**_xyz
        """
        param = self.poses[icam]
        assert p3d.shape[-1] == 3
        R,t = param['R'], param['t']
        p3d_T = p3d.reshape(-1, 3).T
        p3d_algn_T = R@p3d_T + t.reshape(-1,1)
        p3d_algn = p3d_algn_T.T.reshape(*p3d.shape)
        return p3d_algn

    def get_cam_pos_p3d(self) -> np.ndarray:
        """
        Get camera position in 3D.
        """
        cam_pos = np.zeros((len(self.views), 3), dtype=float)
        for i, view in enumerate(self.views):
            param = self.poses[view]
            R,t = param['R'], param['t']
            cam_pos[i] = (-np.linalg.inv(R) @ t.reshape(-1,1)).ravel()
        return cam_pos
    
    def get_cam_direct_p3d(self) -> np.ndarray:
        """
        Get camera direction of X,Y,Z
        """
        cam_oxyz = np.zeros((len(self.views), 4, 3), dtype=float)
        oxyz = np.eye(4)[1:,:] #(3,4)
        for i, view in enumerate(self.views):
            param = self.poses[view]
            R,t = param['R'], param['t']
            cam_oxyz[i] = (np.linalg.inv(R) @ (oxyz-t.reshape(3,1))).T
        cam_xyz = cam_oxyz[:,1:] - cam_oxyz[:,[0]]
        return cam_xyz
    

def build_input_np(views, poses:dict, landmarks:np.ndarray) -> np.ndarray:
    # transform camera poses to numpy array
    n_cameras = len(views)
    poses_list = [[np.squeeze(np.float64(poses[j][key])) for key in ['K', 'R', 't', 'dist']]
                    for j in range(n_cameras)]

    # undistort landmarks
    landmarks_undist_withnan = np.zeros_like(landmarks)
    for j in range(n_cameras):
        K,R,t,dist = poses_list[j]
        points = landmarks[views[j]]
        landmarks_undist_withnan[views[j]] = undistort_points(points, K, dist)

    # triangulate 3D positions from all possible pair of views
    p3d_allview_withnan = []
    for j1, j2 in itertools.combinations(range(n_cameras), 2):
        K1,R1,t1,dist1 = poses_list[j1]
        K2,R2,t2,dist2 = poses_list[j2]
        pts1 = landmarks_undist_withnan[views[j1]]
        pts2 = landmarks_undist_withnan[views[j2]]
        p3d = triangulate(pts1, pts2, K1, R1, t1, None, K2, R2, t2, None)
        p3d_allview_withnan.append(p3d)
    p3d_allview_withnan = np.array(p3d_allview_withnan) #nviewpairs_nsample_xyz

    # get median
    p3d = np.nanmedian(p3d_allview_withnan, axis=0) #nsample_xyz
    return p3d


def project_points(pts, K, R, t, dist=None, image_shape=None):
    #from multiview_calib.singleview_geometry import (undistort_points, project_points)
    pts_ = np.array(pts)
    assert pts_.shape[1]==3
    
    proj = np.dot(K, np.dot(R,pts_.T) + t.reshape(3,1))
    z = proj[2]
    xy = proj[:2].T/z[:,None]
    mask_in_front = z > 0  #original
    # mask_in_front = z > -np.inf #cxf

    rvec = cv2.Rodrigues(R)[0]
    proj = cv2.projectPoints(pts_, rvec, t, K, dist)[0].reshape(-1,2)
    
    if image_shape is not None:
        mask_inside = np.logical_and.reduce([proj[:,0]>0, proj[:,0]<image_shape[1],
                                             proj[:,1]>0, proj[:,1]<image_shape[0]])
        mask_valid = np.logical_and(mask_in_front, mask_inside)
    else:
        mask_valid = mask_in_front    

    return proj, mask_valid


def undistort_points(points, K, distCoeffs, norm_coord=False, newcameramtx=None):
    #from multiview_calib
    points_ = np.reshape(points, (-1,1,2))
    if newcameramtx is None:
        newcameramtx = K
    points_ = cv2.undistortPoints(np.float32(points_), K, distCoeffs, P=newcameramtx, R=None)
    points_ = np.reshape(points_, (-1,2))
    return points_


def triangulate(pts1, pts2, 
                K1=np.eye(3), R1=np.eye(3), t1=np.zeros(3), dist1=None, 
                K2=np.eye(3), R2=np.eye(3), t2=np.zeros(3), dist2=None):
    #from multiview_calib
    pts1_ = np.reshape(pts1, (-1,2)).copy()
    pts2_ = np.reshape(pts2, (-1,2)).copy()  
    
    if dist1 is not None and dist2 is not None:
        pts1_undist = undistort_points(pts1_, K1, dist1)
        pts2_undist = undistort_points(pts2_, K2, dist2)
    else:
        pts1_undist = pts1_
        pts2_undist = pts2_      
        
    P1 = np.dot(K1, np.hstack([R1, t1.reshape(3,1)]))
    P2 = np.dot(K2, np.hstack([R2, t2.reshape(3,1)]))

    tri = cv2.triangulatePoints(P1, P2, pts1_undist.T, pts2_undist.T)
    tri = tri[:3]/tri[3]
    tri = tri.T 
    
    return tri
